package ms.crack.zhilian;


import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import net.sourceforge.tess4j.util.ImageHelper;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.filefilter.TrueFileFilter;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.net.URL;
import java.util.List;

public class ImageORCService {

    public static final int TEST_SIZE = 10;

    private int testSize = TEST_SIZE;
    private String imageUrl = "";
    private String saveDir;
    private String imageType = "gif";

    public ImageORCService(String imageUrl, String saveDir, String imageType) {
        this.imageUrl = imageUrl;
        this.saveDir = saveDir;
        this.imageType = imageType;
    }

    public void doOCR() {
        File dir = new File(saveDir);
        if (!dir.exists()) {
            dir.mkdirs();
        }
        //随机下载图片
        downLoadImge();
        //读取图片
        List<File> files = (List<File>) FileUtils.listFiles(dir, TrueFileFilter.INSTANCE, TrueFileFilter.INSTANCE);
        //OCR解析
        for (int i = 0; i < files.size(); i++) {

            File file = files.get(i);


            ITesseract instance = new Tesseract();  // JNA Interface Mapping
            // ITesseract instance = new Tesseract1(); // JNA Direct Mapping
            try {
                BufferedImage bin = ImageIO.read(file);
                //二值化
                bin = processBinarization(bin);
                String result = instance.doOCR(bin);
                System.out.println("图片:" + file.getAbsolutePath() + " 解析结果:" + result);
                //  重命名文件
//                FileUtils.moveFile(file, new File(file.getParent() + "/" + FilenameUtils.getPrefix(file.getName()) + "_" + result + "." + imageType));
            } catch (TesseractException e) {
                System.err.println(e.getMessage());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

    }


    public static BufferedImage processBinarization(BufferedImage bin) {

        int endX = bin.getWidth();
        int endY = bin.getHeight();
        // 这里对图片黑白处理,增强识别率.这里先通过截图,截取图片中需要识别的部分
        BufferedImage textImage = ImageHelper.convertImageToGrayscale(ImageHelper.getSubImage(bin, 0, 0, endX, endY));
        // 图片锐化,自己使用中影响识别率的主要因素是针式打印机字迹不连贯,所以锐化反而降低识别率
        // textImage = ImageHelper.convertImageToBinary(textImage);
        // 图片放大5倍,增强识别率(很多图片本身无法识别,放大5倍时就可以轻易识,但是考滤到客户电脑配置低,针式打印机打印不连贯的问题,这里就放大5倍)
        textImage = ImageHelper.getScaledInstance(textImage, endX * 5, endY * 5);

        return textImage;
    }


    private void downLoadImge() {
        for (int i = 0; i < testSize; i++) {
            byte[] bytes = downloadUrl(imageUrl);
            File gifFile = new File(saveDir + "/" + (i + 1) + "." + imageType);
            try {
                OutputStream out = new FileOutputStream(gifFile);
                IOUtils.write(bytes, out);
                IOUtils.closeQuietly(out);
            } catch (FileNotFoundException e) {
                e.printStackTrace();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public int getTestSize() {
        return testSize;
    }

    public void setTestSize(int testSize) {
        this.testSize = testSize;
    }


    private byte[] downloadUrl(String toDownload) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

        try {
            byte[] chunk = new byte[4096];
            int bytesRead;
            InputStream stream = (new URL(toDownload)).openStream();

            while ((bytesRead = stream.read(chunk)) > 0) {
                outputStream.write(chunk, 0, bytesRead);
            }

        } catch (IOException e) {
            e.printStackTrace();
            return null;
        }

        return outputStream.toByteArray();
    }
}
